import numpy as np
import scipy.signal as signal

from LTE.lte import LTE,LTEInfo
from LTESync.sequence import PSS,SSS
from LTESync.afc import AFC
from LTE.lte import SSInd,ExtractResources,ofdm_rx_at
class LTESync(LTE):
    def __init__(self,fn,centered=True) -> None:
        LTE.__init__(self,fn,centered)
        self.cell = None
        
    def sync(self):
        peak_value = np.zeros(3)
        peak_index = np.zeros(3, 'int')
        wv = self.waveform[:19200]
        for current_NID2 in np.arange(3, dtype='int'):
            pss = PSS(current_NID2)
            temp = signal.correlate(wv, pss.td.conj(), 'valid')  # correlate over 25 ms
            index = np.argmax(np.abs(temp))
            peak = np.abs(temp[index])
            peak_index[current_NID2] = index
            peak_value[current_NID2] = peak
        self.detected_NID2 = np.argmax(peak_value)
        timingOffset = peak_index[self.detected_NID2]+137
        if timingOffset < 137*4:
            timingOffset += 19200
        id,slot,amp,ok_fdd = self.find_sss(-137,timingOffset)
        if ok_fdd:
            off = timingOffset-960 if slot==0 else timingOffset+9600-960
            off %= 19200
            self.cell = LTEInfo().f(self.re_center).cellid(id).duplex("fdd").offset(off).port(0).set_peak(amp)
        id,slot,amp,ok_tdd = self.find_sss(-(137*3)-1,timingOffset)
        if ok_tdd:
            off = timingOffset-137*3-1-1920 if slot==0 else timingOffset-137*3-1-1920*6
            off %= 19200
            self.cell = LTEInfo().f(self.re_center).cellid(id).duplex("tdd").offset(off).port(0).set_peak(amp)
        return ok_fdd or ok_tdd
    
    def find_sss(self,dt,timingOffset):
        to = timingOffset - 137
        pss = PSS(self.detected_NID2)
        pss_fd = ofdm_rx_at(self.waveform,to,128,4)
        pss_ex = ExtractResources(SSInd(),pss_fd)
        sss_fd = ofdm_rx_at(self.waveform,to+dt,128,4)
        sss_ex = ExtractResources(SSInd(),sss_fd)
        dm_pss = pss_ex*pss.ofd.conj()
        avg = np.mean(dm_pss)
        sss_dm = sss_ex * avg.conj()
        sss_dm /= max(sss_dm.real.max(),sss_dm.imag.max())
        res = np.zeros((168*2,),dtype=np.complex64)
        for id_1 in range(168):
            id_2 = self.detected_NID2
            for slot in range(2):
                inx = id_1*2+slot
                sss = SSS.sss.get(id_1,id_2,slot)
                res[inx] = np.vdot(sss_dm,sss.ofd)
        d = np.argmax(np.abs(res))
        peak = np.abs(res[d])
        nd = np.std(res)
        # print(self.re_center,d,peak,nd)
        return d//2*3+self.detected_NID2,d%2,peak,(peak>4*nd and peak>10)
        
        
    def afc_fdd(self,center_f):
        return AFC.afc(self.waveform,center_f)
    
    def afc(self):
        if self.cell is None:
            print("No cell Detected")
            return 0,0
        r = []
        for pid in range(2):
            symbol = AFC.refs[self.cell.cid].symbol_of_frame(pid,self.waveform,self.cell.off)
            f = np.abs(np.fft.fft(symbol))
            df = f.argmax()
            amp = f[df]
            r.append((amp,pid,df))
        if r[0][0]>r[1][0]:
            self.cell.port(0)
            df = r[0][2]
        else:
            self.cell.port(1)
            df = r[0][2]
        # print(r)
        s = np.zeros(19200*100,dtype=complex)
        n = len(self.waveform) // 19200 - 1
        for f in range(n):
            s[f*19200:f*19200+19200] = AFC.refs[self.cell.cid].symbol_of_frame(
                self.cell.pid,self.waveform[f*19200:],self.cell.off)
        point = s.shape[0]
        fc = np.fft.fftshift(np.fft.fft(s))
        dc_fc = fc[point//2-5000:point//2+5000]
        df = (np.abs(dc_fc).argmax()-5000)
        return df,dc_fc[df+5000]